package cn.monkey.family.data;

import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

public interface TreeMapper<K,T,R> {

    default  Collection<T> getRoot(Collection<T> collection,
                                         Map<K, T> ktMap,
                                         Function<T, K> parentKeyFunction) {
        List<T> root = new ArrayList<>();
        for (T t : collection) {
            K parentKey = parentKeyFunction.apply(t);
            if (parentKey == null) {
                root.add(t);
                continue;
            }
            if (!ktMap.containsKey(parentKey)) {
                root.add(t);
            }
        }
        return root;
    }


    default Collection<R> toTree(Collection<T> collection,
                                           Function<T, R> tToRFunc,
                                           Function<T, K> keyFunction,
                                           Function<T, K> parentKeyFunction,
                                           BiConsumer<R, Collection<R>> setChildFunc,
                                           int maxDeep) {
        if (CollectionUtils.isEmpty(collection)) {
            return Collections.emptyList();
        }

        Map<K, T> ktMap = collection.stream().collect(Collectors.toMap(keyFunction, Function.identity()));
        Collection<T> root = this.getRoot(collection, ktMap, parentKeyFunction);
        if (maxDeep == 0) {
            return root.stream().map(tToRFunc).collect(Collectors.toList());
        }
        if (CollectionUtils.isEmpty(root)) {
            throw new IllegalArgumentException("can not find root, maybe circle relation collection?");
        }
        Map<K, List<T>> parentIdKeyTMap = collection.stream()
                .filter(t -> parentKeyFunction.apply(t) != null)
                .collect(Collectors.groupingBy(parentKeyFunction));
        Collection<R> rs = new ArrayList<>();
        for (T t : root) {
            K key = keyFunction.apply(t);
            List<T> ts = parentIdKeyTMap.get(key);
            R r = tToRFunc.apply(t);
            rs.add(r);
            if (CollectionUtils.isEmpty(ts)) {
                continue;
            }
            setChildFunc.accept(r, this.toTree(ts, tToRFunc, keyFunction, parentIdKeyTMap, setChildFunc, 1, maxDeep));
        }
        return rs;
    }

    default  Collection<R> toTree(Collection<T> collection,
                                           Function<T, R> tToRFunc,
                                           Function<T, K> keyFunction,
                                           Map<K, List<T>> parentKeyTMap,
                                           BiConsumer<R, Collection<R>> setChildFunc,
                                           int currentDeep,
                                           int maxDeep) {
        if (CollectionUtils.isEmpty(collection)) {
            return Collections.emptyList();
        }
        if (maxDeep >= 0 && currentDeep >= maxDeep) {
            return collection.stream().map(tToRFunc).collect(Collectors.toList());
        }
        List<R> rs = new ArrayList<>(collection.size());
        for (T t : collection) {
            K key = keyFunction.apply(t);
            List<T> ts = parentKeyTMap.get(key);
            R r = tToRFunc.apply(t);
            rs.add(r);
            if (CollectionUtils.isEmpty(ts)) {
                continue;
            }
            setChildFunc.accept(r, this.toTree(ts, tToRFunc, keyFunction, parentKeyTMap, setChildFunc, currentDeep + 1, maxDeep));
        }
        return rs;
    }

}
